import json
from .tasks import save_private_message
from channels.generic.websocket import AsyncWebsocketConsumer
from django.contrib.auth import get_user_model
from .models import get_room_key
import channels
from .models import PrivateRoom
from django.utils import timezone
import logging
from utils.async_redis import get_redis_pool
from asgiref.sync import sync_to_async
from channels.db import database_sync_to_async
User = get_user_model()
logger = logging.getLogger(__name__)

class PrivateChat(AsyncWebsocketConsumer):
    # 只有进入私聊页面才触发这个consumer
    def __init__(self):
        super().__init__(self)
        self.room_key = None
        self.user = None
        self.private_group_name = None
        self.private_room = None
        self.peer_id = None
        self.peer_user = None
    async def connect(self):
        self.room_key = self.scope['url_route']['kwargs']['room_key']
        self.peer_id = self.scope['url_route']['kwargs']['peer_id']
        self.peer_user = await self.get_user()
        self.user = self.scope['user']
        if not self.user.is_authenticated:
            await self.close(code=4001)  # 用户未登录
            return
        if not await self.is_room_exit() or not self.peer_user:
            await self.close(code=4002)  # 目标私聊不存在
            return
        self.private_group_name = f'private_{min(self.user.id, self.peer_id)}_{max(self.user.id, self.peer_id)}'
        await self.channel_layer.group_add(self.private_group_name, self.channel_name)
        redis = await get_redis_pool()
        await redis.sadd(f"user_private_rooms:{self.user.id}", self.private_group_name)
        # 把当前连接加入到redis中，方便查询 sadd就是将当前私聊组名加入到用户id集合中，无重复
        await self.accept()
    async def receive(self, text_data=None, bytes_data=None):
        try:
            data = json.loads(text_data)
            msg = data.get('message')
            if not msg or not msg.strip():
                return
            event = {
                'type':'private_message',
                'message':msg,
                'sender_name':self.user.username,
                'timestamp': timezone.now().isoformat(),
                'is_temp': True,
            }
            await self.channel_layer.group_send(
                self.private_group_name,
                event
            )
            save_private_message.delay(
                key=self.room_key,
                sender_id=self.user.id,
                receiver_id=self.peer_user.id,
                content=msg,
                timestamp=event['timestamp']
            )
            in_room = await self.check_user_in_room(self.peer_id, self.private_group_name)
            if not in_room:
                # 给接收者专属提醒组发提醒
                await self.channel_layer.group_send(
                    f'notice_{self.peer_id}',
                    {
                        'type': 'notify_func',
                        'message': f'你有一条来自用户{self.user.username}的新消息',
                    }
                )
        except json.JSONDecodeError:
            logger.error("Invalid JSON received")
        except Exception as e:
            logger.exception("Error processing message")
    async def private_message(self,event):
        await self.send(text_data=json.dumps({
            'type': 'private_message',
            'message':event['message'],
            'sender_name':event['sender_name'],
            'timestamp':event['timestamp'],
            'is_temp': event['is_temp'],
        }
        ))

    async def disconnect(self, code):
        if self.private_group_name:
            await self.channel_layer.group_discard(self.private_group_name,self.channel_name)
            redis = await get_redis_pool()
            await redis.srem(f"user_private_rooms:{self.user.id}", self.private_group_name)
            logger.info(f"Private chat ended between {self.user.id} and {self.peer_id} with code {code}")
        else:
            raise ValueError("断开连接出错")
    @staticmethod
    async def check_user_in_room(peer_id, private_group_name):
        # 不依赖self 可以是静态
        redis = await get_redis_pool()
        ## 从 Redis 里取出用户加入的所有私聊组集合
        rooms = await redis.smembers(f"user_private_rooms:{peer_id}")
        return private_group_name in rooms

    @database_sync_to_async
    def is_room_exit(self):
        return PrivateRoom.objects.filter(key=self.room_key).exists()

    @database_sync_to_async
    def get_user(self):
        return User.objects.filter(id=self.peer_id).first()

class PrivateNotice(AsyncWebsocketConsumer):
    # 不在私聊页面也触发这个consumer，主要是别人在私聊页面使用这个consumer里的函数进行触发
    def __init__(self):
        super().__init__(self)
        self.user = None
        self.group_id = None
    async def connect(self):
        self.user = self.scope['user']
        if not self.user.is_authenticated:
            await self.close(code=4001) # 用户未登录,不接收提醒
            return
        self.group_id = f'notice_{self.user.id}'
        await self.channel_layer.group_add(self.group_id,self.channel_name)
        await self.accept()
    async def disconnect(self, close_code):
        await self.channel_layer.group_discard(self.group_id, self.channel_name)
        logger.info(f"Private notice{self.user.id} ended with code {close_code}")
    async def notify_func(self,event):
        print('发送了提醒')
        await self.send(text_data=json.dumps(
            {'type': 'notify_func',
             'message':event['message'],
             }
        ))
# chat/consumers.py
import json
from channels.generic.websocket import AsyncWebsocketConsumer
from server.llm_agent import AgentService
from User.models import CustomerUser

class DeepSeekChatConsumer(AsyncWebsocketConsumer):
    @database_sync_to_async
    def get_own_user(self):
        return User.objects.filter(id=self.user.id).first()
    async def connect(self):
        self.user = self.scope["user"]

        await self.accept()

    async def receive(self, text_data):
        data = json.loads(text_data)
        message = data.get("message")
        # target_user_id = data.get("target_user_id")

        # 保存用户消息
        # ChatRecord.objects.create(
        #     user_id=self.user.id,
        #     target_user_id=target_user_id,
        #     content=message
        # )

        # 获取用户 Agent

        profile = await self.get_own_user()
        info = {'id': self.user.id, 'api': profile.deepseekapi}
        agent = AgentService.get_agent(info)

        # Agent 自主判断是否调用工具
        from asgiref.sync import async_to_sync

        reply_text = await sync_to_async(agent.invoke)({"input": message})
        # reply_text = agent.run(message)

        # 保存 Agent 回复
        # ChatReply.objects.create(
        #     user_id=self.user.id,
        #     target_user_id=target_user_id,
        #     content=reply_text
        # )

        # 返回前端
        await self.send(text_data=json.dumps({
            "type": "agent_reply",
            "reply": reply_text
        }))

